import torch
from torch import nn
from torch.nn import functional as F


IMG_SIZE = 224
LR = 0.005


class Residual(nn.Module):
    def __init__(self, in_channels, out_channels, use_conv1x1=False, stride=1):
        super(Residual, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_channels)
        )
        self.shortcut = nn.Sequential()
        if (use_conv1x1):
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride),
                nn.BatchNorm2d(out_channels)
            )


    def forward(self, input):
        y = self.net(input)
        y_shortcut = self.shortcut(input)
        return F.relu(y + y_shortcut)


class ResNet18(nn.Module):
    """ResNet 18"""
    def __init__(self, num_classes, in_channels=3):
        super(ResNet18, self).__init__()
        self.features = nn.Sequential(
            ## 3x3 conv, maxpool
            nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),

            ## 3x3 conv, 64
            Residual(64, 64),
            Residual(64, 64),

            ## 3x3 conv, 128
            Residual(64, 128, use_conv1x1=True, stride=2),
            Residual(128, 128),

            ## 3x3 conv, 256
            Residual(128, 256, use_conv1x1=True, stride=2),
            Residual(256, 256),

            ## 3x3 conv, 512
            Residual(256, 512, use_conv1x1=True, stride=2),
            Residual(512, 512),
        )

        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(512, num_classes)
        )
        


    def forward(self, input):
        y = self.features(input)
        return self.classifier(y)


    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR


class ResNet34(nn.Module):
    """ResNet 34"""
    def __init__(self, num_classes, in_channels=3):
        super(ResNet34, self).__init__()
        self.features = nn.Sequential(
            ## 3x3 conv, maxpool
            nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),

            ## 3x3 conv, 64
            Residual(64, 64),
            Residual(64, 64),
            Residual(64, 64),

            ## 3x3 conv, 128
            Residual(64, 128, use_conv1x1=True, stride=2),
            Residual(128, 128),
            Residual(128, 128),
            Residual(128, 128),

            ## 3x3 conv, 256
            Residual(128, 256, use_conv1x1=True, stride=2),
            Residual(256, 256),
            Residual(256, 256),
            Residual(256, 256),
            Residual(256, 256),
            Residual(256, 256),

            ## 3x3 conv, 512
            Residual(256, 512, use_conv1x1=True, stride=2),
            Residual(512, 512),
            Residual(512, 512),
        )

        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(512, num_classes)
        )
        


    def forward(self, input):
        y = self.features(input)
        return self.classifier(y)


    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR


class Bottleneck(nn.Module):
    def __init__(self, in_channels, mid_channels, out_channels, use_conv1x1=False, stride=1):
        super(Bottleneck, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(),
            nn.Conv2d(mid_channels, mid_channels, kernel_size=3, stride=stride, padding=1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(),
            nn.Conv2d(mid_channels, out_channels, kernel_size=1),
            nn.BatchNorm2d(out_channels),
        )
        self.shortcut = nn.Sequential()
        if (use_conv1x1):
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride),
                nn.BatchNorm2d(out_channels)
            )


    def forward(self, input):
        y = self.net(input)
        y_shortcut = self.shortcut(input)
        return F.relu(y + y_shortcut) 



class ResNet50(nn.Module):
    """ResNet 50"""
    def __init__(self, num_classes, in_channels=3):
        super(ResNet50, self).__init__()
        self.features = nn.Sequential(
            ## Stage0 3x3 conv, maxpool
            nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),

            ## Stage1 3x3 conv, 256
            Bottleneck(64, 64, 256, use_conv1x1=True, stride=1),
            Bottleneck(256, 256, 256),
            Bottleneck(256, 256, 256),

            ## Stage2 3x3 conv, 512
            Bottleneck(256, 128, 512, use_conv1x1=True, stride=2),
            Bottleneck(512, 128, 512),
            Bottleneck(512, 128, 512),
            Bottleneck(512, 128, 512),

            ## 3x3 conv, 1024
            Bottleneck(512, 256, 1024, use_conv1x1=True, stride=2),
            Bottleneck(1024, 256, 1024),
            Bottleneck(1024, 256, 1024),
            Bottleneck(1024, 256, 1024),
            Bottleneck(1024, 256, 1024),
            Bottleneck(1024, 256, 1024),

            ## 3x3 conv, 2048 
            Bottleneck(1024, 512, 2048, use_conv1x1=True, stride=2),
            Bottleneck(2048, 512, 2048),
            Bottleneck(2048, 512, 2048)
        )

        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(2048, num_classes)
        )
        


    def forward(self, input):
        y = self.features(input)
        return self.classifier(y)


    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR


class ResNet101(nn.Module):
    """ResNet 101"""
    def __init__(self, num_classes, in_channels=3):
        super(ResNet101, self).__init__()
        
        conv4 = []
        for _ in range(22):
            conv4.append(Bottleneck(1024, 256, 1024))

        self.features = nn.Sequential(
            ## Stage0 3x3 conv, maxpool
            nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),

            ## Stage1 3x3 conv, 256
            Bottleneck(64, 64, 256, use_conv1x1=True, stride=1),
            Bottleneck(256, 256, 256),
            Bottleneck(256, 256, 256),

            ## Stage2 3x3 conv, 512
            Bottleneck(256, 128, 512, use_conv1x1=True, stride=2),
            Bottleneck(512, 128, 512),
            Bottleneck(512, 128, 512),
            Bottleneck(512, 128, 512),

            ## 3x3 conv, 1024
            Bottleneck(512, 256, 1024, use_conv1x1=True, stride=2),
            *conv4,

            ## 3x3 conv, 2048 
            Bottleneck(1024, 512, 2048, use_conv1x1=True, stride=2),
            Bottleneck(2048, 512, 2048),
            Bottleneck(2048, 512, 2048)
        )

        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(2048, num_classes)
        )
        


    def forward(self, input):
        y = self.features(input)
        return self.classifier(y)


    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR


class ResNet152(nn.Module):
    """ResNet 152"""
    def __init__(self, num_classes, in_channels=3):
        super(ResNet152, self).__init__()
        
        conv3, conv4 = [], []
        for _ in range(7):
            conv3.append(Bottleneck(256, 256, 256))
        for _ in range(35):
            conv4.append(Bottleneck(1024, 256, 1024))

        self.features = nn.Sequential(
            ## Stage0 3x3 conv, maxpool
            nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),

            ## Stage1 3x3 conv, 256
            Bottleneck(64, 64, 256, use_conv1x1=True, stride=1),
            Bottleneck(256, 256, 256),
            Bottleneck(256, 256, 256),

            ## Stage2 3x3 conv, 512
            Bottleneck(256, 128, 512, use_conv1x1=True, stride=2),
            *conv3,    # x7

            ## 3x3 conv, 1024
            Bottleneck(512, 256, 1024, use_conv1x1=True, stride=2),
            *conv4,    # x35

            ## 3x3 conv, 2048 
            Bottleneck(1024, 512, 2048, use_conv1x1=True, stride=2),
            Bottleneck(2048, 512, 2048),
            Bottleneck(2048, 512, 2048)
        )

        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(2048, num_classes)
        )
        


    def forward(self, input):
        y = self.features(input)
        return self.classifier(y)


    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR